'''
Created on Nov 23, 2010

@author: oabalbin
'''
import os
import subprocess
#import pysam
import exome.gatk_cluster.picard_commands_cluster as mypicard
import exome.gatk_cluster.samtools_commands_cluster as mysam
import exome.gatk_cluster.cluster_jobs_header as jh
from collections import deque



def count_covariates(reference_genome, dbsnp_ref, bam_file_name, outfile, use_mem, num_cores, path_to_gatk):
    '''
    Generated the file of count covariates, which is required for base quality recalibration
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/GenomeAnalysisTK.jar 
    -l INFO -nt 6 -R /exds/projects/alignment_indexes/gatk/hg19/hg19.fa 
    -B:mask,VCF /exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf 
    -I s_3_12_sequence.hg19.aln2.rmdup.sorted.bam -T CountCovariates -cov ReadGroupCovariate 
    -cov QualityScoreCovariate -cov CycleCovariate -cov DinucCovariate 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv
    
    
    Note: you can use the --process_nth_locus/ -pN 3 to speed up the process. 
    Which force the program to evaluate every 3 loci, which makes to skip in general ~66 % data
    when dbsnp is also included. Important for Large BAM files
    '''
    
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    #outfile=bam_file_name.replace('.bam','.recal_data.csv')
    args = ['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,
            '-l', 'INFO', '-nt', str(num_cores), '-R',reference_genome, 
            '-B:dbsnp,VCF', dbsnp_ref, '-I', bam_file_name, '-T','CountCovariates',
            '-cov', 'ReadGroupCovariate','-cov', 'QualityScoreCovariate', '-cov', 
            'CycleCovariate', '-cov', 'DinucCovariate', '-recalFile',outfile,
            '-U', 'ALLOW_UNINDEXED_BAM']
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def table_recalibration(reference_genome,bam_file_name,recal_file_name, outfile, use_mem, path_to_gatk):
    '''
    Generates a new bam file with recalibrated quality scores
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/GenomeAnalysisTK.jar 
    -l INFO -R /exds/projects/alignment_indexes/gatk/hg19/hg19.fa -I s_3_12_sequence.hg19.aln2.rmdup.sorted.bam 
    -T TableRecalibration -o s_3_12_sequence.aln2.rmdup.sorted.recal.bam 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.before.recal_data.csv
    The original function in gatk does not support multiple threads yet
    '''
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    #outfile=bam_file_name.replace('.bam','.recal.bam')
    
    args=['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,
            '-l', 'INFO','-R',reference_genome, '-I', bam_file_name, '-T','TableRecalibration',
            '-o',outfile, '-recalFile',recal_file_name,
            '-U', 'ALLOW_UNINDEXED_BAM']
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def analyze_covariates(recal_file_name,outputdir, resources_folder, rscipt_path, use_mem, path_to_gatk):
    '''
    In this script is important to specify the path to resources in which you have the Rscripts for making the plots and the -R scripts
    -resources /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources (resources_folder)
    -Rscript /exds/sw/local/R-project/bin/Rscript (rscipt_path)
    -ignoreQ 5
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/AnalyzeCovariates.jar 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv 
    -outputDir /exds/users/oabalbin/projects/snps/exomes/aM18/recal_analysis/before/ 
    -resources /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/ 
    -Rscript /exds/sw/local/R-project/bin/Rscript
    '''
    
    gatk_command=path_to_gatk+'AnalyzeCovariates.jar'
    args=['java','-Xmx'+str(use_mem)+'m', '-jar',gatk_command,
            '-recalFile',recal_file_name,'-outputDir',outputdir,'-resources',resources_folder,'-Rscript',rscipt_path]

    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def check_create_dir(root_path, dir_name):
    '''
        if not os.path.isdir(full_path_name):
            os.mkdir( full_path_name )
    '''
    if not os.path.isdir(root_path):
            os.mkdir( root_path )
            
    subfolder=os.path.join(root_path,dir_name)
    
    if not os.path.isdir(subfolder):
            os.mkdir( subfolder )

    '''
    args = ['mkdir', full_path_name]
    command = ",".join(args).replace(',',' ').replace(';',',')
    '''
 
    return subfolder



def main_baseQ_recalibration(ref_genome, ref_dbsnp, bam_file_name, thislane, core_mem, num_cores, 
                             configrun, recal_analysis_outputdir, previous_process_jobid, 
                             jname, my_email, jobrunfunc):
    '''
    This is a wrap for the base quality recalibration using  gatk (Broad).
    Input: reference genome, vcf dbsnp database, indexed original bam file, 
        memory available, num of cores to use, analysis output directory
        path to gatk, resources folder, and RScript
    
    Returns: recalibrate bam file path, Writes recalibrated bam file and 
        Count covariate tables
        
    gatk quality recalibration steps
    CountCovariates, TableRecalibrate, 
    samtools index on the recalibrated bam file, CountCovariates
    analyze covariates before and after recalibration
    Most of the parameters in this function should go into a dictionary structure.
    '''
    '''    
    path_to_gatk, resources_folder,rscipt_path, path_to_sam =\
    gatk_run_dict['path_to_gatk'], gatk_run_dict['resources_folder'],\
    gatk_run_dict['rscipt_path'], gatk_run_dict['path_to_sam']
    '''
    
    path_to_gatk, path_to_picard, path_to_sam = configrun.gatk_path, configrun.picard_path, \
                                                configrun.samtools_path
    resources_folder,rscipt_path = configrun.gatk_resource_path, configrun.rscript_bin

    
    # Generate count covariates table for recalibration   
    node_memory = 45000.0
    node_processors = 12
    single_core=1
    single_mem = int(float(node_memory) / node_processors)

    num_cores = num_cores/2
    core_mem = single_mem*num_cores
    extra_mem=16000
    
    wt_short='24:00:00'
    wt_count_covariates="60:00:00"

    pre_recal_table_file = thislane.pre_recal_table_file
    command = count_covariates(ref_genome, ref_dbsnp, bam_file_name,
                               pre_recal_table_file, core_mem, num_cores, path_to_gatk)
    
    jobidcv = jobrunfunc(jname+'cv', command, num_cores, cwd=None, walltime=wt_count_covariates, pmem=None, 
                              deps=previous_process_jobid, stdout=None, email_addresses=my_email) # deps=None just for test this job depends on  previous_process_jobid       
        
    # Recalibrate original bam file
    # Package 1 starts: into one script
    recal_bam_file_name=thislane.recal_bam_file_name
    command = table_recalibration(ref_genome, bam_file_name,
                                              pre_recal_table_file, recal_bam_file_name, 
                                              extra_mem, path_to_gatk)
    
    jobidtr = jobrunfunc(jname+'tr', command, single_core, cwd=None, walltime=wt_short, pmem=extra_mem, 
                              deps=jobidcv, stdout=None, email_addresses=my_email)            

    
    # Index recalibrated bam file
    recal_bam_indexed_name=thislane.recal_bam_indexed_name
    command  = mypicard.sortIndexSam(recal_bam_file_name, recal_bam_indexed_name, 
                                     single_mem, path_to_picard)
    
    #command, recal_bam_indexed_name = mysam.index_bam_file(recal_bam_file_name,path_to_sam)
    jobidib = jobrunfunc(jname+'ib', command, single_core, cwd=None, walltime=wt_short, pmem=None, 
                              deps=jobidtr, stdout=None, email_addresses=my_email)   #jobidtr         

    
    # Package 1 ends:
    
    # Generate count covariates table for recalibrated bam file
    # Here maybe using more cores
    post_recal_table_file =thislane.post_recal_table_file 
    command = count_covariates(ref_genome, ref_dbsnp, recal_bam_indexed_name, 
                               post_recal_table_file, core_mem, num_cores, path_to_gatk)
    
    jobidcv2 = jobrunfunc(jname+'cv2', command, num_cores, cwd=None, walltime=wt_count_covariates, pmem=None, 
                              deps=jobidib, stdout=None, email_addresses=my_email)
    
    # Analyze covariates before and after recalibration
    #Creating folder should be done for another job during the setup
        
    before_recal_results_path = check_create_dir(recal_analysis_outputdir,'before')
    after_recal_results_path = check_create_dir(recal_analysis_outputdir,'after')    
    # use the retcode to print to a logfile
    commandB = analyze_covariates(pre_recal_table_file, before_recal_results_path, 
                                 resources_folder, rscipt_path, single_mem, path_to_gatk)
    #
    commandA = analyze_covariates(post_recal_table_file, after_recal_results_path, 
                                  resources_folder, rscipt_path, single_mem, path_to_gatk)
    
    # Change this line
    command = commandA+'\n'+commandB
    
    jobidab = jobrunfunc(jname+'ab', command, single_core, cwd=None, walltime=wt_short, pmem=None, 
                              deps=jobidcv2, stdout=None, email_addresses=my_email) #deps=jobidcv2
    '''
    jobidab = jobrunfunc(jname+'ab', command, single_core, cwd=None, walltime=wt_short, pmem=None, 
                              deps=previous_process_jobid, stdout=None, email_addresses=my_email) #deps=jobidcv2
    '''

            
    return [jobidab], recal_bam_indexed_name

    
## Test Script
if __name__ == '__main__':
    '''
    ref_genome='/exds/projects/alignment_indexes/gatk/hg19/hg19.fa'
    ref_dbsnp='/exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf'
    bam_file_name='/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.bam'
    use_mem=24
    num_cores=6
    recal_analysis_outputdir='/exds/users/oabalbin/projects/snps/exomes/aM18/test/recal_analysis/'
     
    gatk_run_dict = {'path_to_gatk':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/', 
                     'resources_folder':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/', 
                     'rscipt_path':'/exds/sw/local/R-project/bin/Rscript'}
    '''
    gatk_run_dict = {'path_to_gatk':'/nobackup/med-mctp/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/',
                 'path_to_picard':'/nobackup/med-mctp/sw/bioinfo/picard/picard-tools-1.35/', 
                 'path_to_sam':'/nobackup/med-mctp/sw/bioinfo/samtools/samtools-0.1.10/',
                 'resources_folder':'/nobackup/med-mctp/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/', 
                 'rscipt_path':'/home/software/rhel5/R/2.10.1-gcc/bin/Rscript',
                 'use_mem':8000, 'num_cores':1,
                 'ref_genome':'/nobackup/med-mctp/sw/alignment_indexes/gatk/hg19/hg19.fa', 
                 'snpdb_file':'/nobackup/med-mctp/sw/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed.vcf',
                 'indeldb_file':'/nobackup/med-mctp/sw/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed.vcf',
                 'path_to_intervals':'/nobackup/med-mctp/oabalbin/test/',
                 'recal_analysis_outputdir':'/nobackup/med-mctp/oabalbin/test/recal_analysis/',
                 'temp_dir':'/nobackup/med-mctp/oabalbin/test/temp/',
                 'qsubfile':'/nobackup/med-mctp/oabalbin/test/',
                 'out_dir':'/nobackup/med-mctp/oabalbin/test/',
                 'email':['alebalbin@gmail.com']    
                 }

    bam_file_name = '/nobackup/med-mctp/oabalbin/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.realigned.fixMate.markdup.bam'
    core_mem, num_cores = 45000.0, 12
    previous_process_jobid='go'
    jname, my_email = 'aMqc',gatk_run_dict['email']
    main_baseQ_recalibration(gatk_run_dict['ref_genome'], gatk_run_dict['snpdb_file'], bam_file_name, core_mem, num_cores, 
                             gatk_run_dict, gatk_run_dict['recal_analysis_outputdir'], previous_process_jobid, jname, my_email)
    
    
    
    




    
    
    
    
